#pragma once

#include <vector>
#include <string>
#include <fstream>
#include <iostream>
#include <cmath>
#include <algorithm>
#include <cstdlib>
#include <ctime>
#include <mpi.h>

// Debug log levels
#define DEBUG_NONE 0
#define DEBUG_ERROR 1
#define DEBUG_WARNING 2
#define DEBUG_INFO 3
#define DEBUG_DETAIL 4
#define DEBUG_VERBOSE 5

// Debug logging macro
#define DEBUG_LOG(level, msg) \
    if (debug_level >= level) { \
        std::cout << msg << std::endl; \
    }

namespace TensorMatMultiCavity {

// Global MPI variables
extern int mpi_rank;
extern int mpi_size;
extern int mpi_procs_x;
extern int mpi_procs_y;
extern int mpi_coords[2];
extern int debug_level;

// Forward declarations
class MultiCavityTensorMatrix;
class MultiCavityGridParams;

// Function declarations for grid operators
void generate_sem_operators_1d(int N, int Ncell, double L,
                             std::vector<double>& points,
                             std::vector<std::vector<double>>& D1,
                             std::vector<std::vector<double>>& D2);

void generate_chebyshev_operators(int N, double L,
                                std::vector<double>& points,
                                std::vector<std::vector<double>>& D1,
                                std::vector<std::vector<double>>& D2);

// Structure to hold local grid information
struct LocalGridInfo {
    // Global dimensions
    int global_imt;
    int global_jmt;
    int global_kmt;

    // Local dimensions
    int local_imt;
    int local_jmt;
    int local_kmt;

    // Local grid starting indices
    int local_ims;
    int local_jms;
    int local_kms;

    // Local grid ending indices
    int local_ime;
    int local_jme;
    int local_kme;

    // MPI information
    int rank;
    int size;
    int offset;
    int halo_width;

    // Neighbor ranks
    int north_rank;
    int south_rank;
    int east_rank;
    int west_rank;

    // Total number of points (including padding)
    int total_padded_points;

    // Added for communication purposes
    int mpi_rank;

    // Constructor
    LocalGridInfo() : global_imt(0), global_jmt(0), global_kmt(0),
                     local_imt(0), local_jmt(0), local_kmt(0),
                     local_ims(0), local_jms(0), local_kms(0),
                     local_ime(0), local_jme(0), local_kme(0),
                     rank(0), size(0), offset(0), halo_width(1),
                     north_rank(MPI_PROC_NULL), south_rank(MPI_PROC_NULL),
                     east_rank(MPI_PROC_NULL), west_rank(MPI_PROC_NULL),
                     total_padded_points(0), mpi_rank(0) {}

    // Print function for debugging
    void print() const {
        std::cout << "Local Grid Info:" << std::endl;
        std::cout << "  Global dimensions: " << global_imt << "x" << global_jmt << "x" << global_kmt << std::endl;
        std::cout << "  Local dimensions: " << local_imt << "x" << local_jmt << "x" << local_kmt << std::endl;
        std::cout << "  Local indices: [" << local_ims << ":" << local_ime << "]x[" 
                  << local_jms << ":" << local_jme << "]x[" << local_kms << ":" << local_kme << "]" << std::endl;
        std::cout << "  MPI rank: " << rank << "/" << size << std::endl;
        std::cout << "  Neighbors: N=" << north_rank << " S=" << south_rank 
                  << " E=" << east_rank << " W=" << west_rank << std::endl;
        std::cout << "  Total padded points: " << total_padded_points << std::endl;
    }
};

// Structure to hold grid parameters
struct GridParams {
    int imt;               // Grid dimensions
    int jmt;
    int kmt;
    bool is_global;        // Whether this is a global grid
    int cavity_id;         // ID of the cavity this grid belongs to
    
    // Helper function to calculate padded index
    int padded_index(int i, int j, int k) const {
        return (k * jmt + j) * imt + i;
    }
};

// Structure to represent a connection between two cavities
struct CavityConnection {
    enum class Type {
        FACE_X,
        FACE_Y,
        FACE_Z
    };

    int cavity1_id;
    int cavity2_id;
    Type type;
    int face1_index;
    int face2_index;

    CavityConnection(int c1, int c2, Type t, int f1, int f2)
        : cavity1_id(c1), cavity2_id(c2), type(t), face1_index(f1), face2_index(f2) {}
};

// Structure to represent a cavity in the grid
struct Cavity {
    LocalGridInfo local_info;
    int data_offset;

    Cavity(const LocalGridInfo& info, int offset)
        : local_info(info), data_offset(offset) {}
};

// Class to manage multiple cavity grid parameters
class MultiCavityGridParams {
public:
    // Constructor
    MultiCavityGridParams() : num_cavities(0), total_points(0) {}

    // Add a cavity to the grid
    void add_cavity(const LocalGridInfo& info) {
        int offset = total_points;
        cavities.emplace_back(info, offset);
        total_points += info.total_padded_points;
        num_cavities++;
    }

    // Add a connection between cavities
    void add_connection(int cavity1_id, int cavity2_id, CavityConnection::Type type,
                       int face1_index, int face2_index) {
        connections.emplace_back(cavity1_id, cavity2_id, type, face1_index, face2_index);
    }

    // Getters
    int get_num_cavities() const { return num_cavities; }
    int get_total_points() const { return total_points; }
    const std::vector<Cavity>& get_cavities() const { return cavities; }
    const std::vector<CavityConnection>& get_connections() const { return connections; }

    // Make these members public for direct access
    std::vector<Cavity> cavities;
    std::vector<CavityConnection> connections;

private:
    int num_cavities;
    int total_points;
};

// Function declarations for grid decomposition
void decompose_grid(int global_imt, int global_jmt, int global_kmt, LocalGridInfo& local_grid);
void decompose_grid_global_mode(int global_imt, int global_jmt, int global_kmt, LocalGridInfo& local_grid);
void communicate_local_boundaries(float* vec, const LocalGridInfo& local_grid);
void communicate_global_boundaries(float* vec, const LocalGridInfo& local_grid);
void communicate_boundaries_multi_cavity(float* data, const MultiCavityGridParams& grid);

// Function to exchange data between connected cavities
void exchange_cavity_data(float* data1, float* data2, 
                        const LocalGridInfo& info1, const LocalGridInfo& info2,
                        const CavityConnection& connection);

// Function to communicate boundaries for a single cavity
void communicate_boundaries(float* data, const LocalGridInfo& info);

// Helper function declarations
void jacobi_gauss_quadrature(int N, double alpha, double beta, std::vector<double>& points, std::vector<double>& weights);
double legendre_polynomial(int n, double x);
void eigen_tridiagonal(const std::vector<double>& alpha, const std::vector<double>& beta,
                      std::vector<double>& eigenvalues, std::vector<double>& eigenvectors);
void legendre_gauss_lobatto(int N, std::vector<double>& points, std::vector<double>& weights);

// Utility functions
float dot_product_multi_cavity(const float* x, const float* y, const MultiCavityGridParams& grid);
float norm_squared_multi_cavity(const float* x, const MultiCavityGridParams& grid);
void dump_matrix_to_file(const MultiCavityTensorMatrix& mat, const std::string& filename);
void compare_matrices(const MultiCavityTensorMatrix& mat1, const MultiCavityTensorMatrix& mat2);

// Multi-cavity tensor matrix class
class MultiCavityTensorMatrix {
public:
    // Constructor
    MultiCavityTensorMatrix(int imt, int jmt, int kmt);
    
    // Destructor
    ~MultiCavityTensorMatrix();
    
    // Access element
    float& operator()(int i, int j, int k);
    float operator()(int i, int j, int k) const;
    
    // Initialize with test pattern
    void initialize_test_pattern();
    
    // Initialize with random values
    void initialize_random();
    
    // Matrix-vector multiplication
    void mat_vec_mult(const float* vec, float* result) const;
    
    // Convert to CSR format
    void convert_to_csr(std::vector<float>& values,
                       std::vector<int>& column_indices,
                       std::vector<int>& row_pointers) const;
    
    // Print matrix info
    void print_info() const;
    
    // Debug: verify connections
    void verify_connections() const;
    
    // Debug: communicate boundaries
    void communicate_boundaries();
    
    // Member variables
    int imt, jmt, kmt;  // Matrix dimensions
    float* data;  // Matrix data
};

} // namespace TensorMatMultiCavity 